package ratelimit

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	commonredis "gitee.com/worklz/yunj-blog-go-micro-common/redis"
	"gitee.com/worklz/yunj-blog-go-micro-common/util"
	"github.com/redis/go-redis/v9"
	"strings"
	"sync"
	"time"

	"golang.org/x/time/rate"
)

// Rule 定义限流规则
type Rule struct {
	Name        string        // 规则名称
	Size        time.Duration // 滑动窗口大小
	MaxRequests int           // 窗口内最大请求数
	LocalRule   LocalRule     // 本地限流器配置
}

// LocalRule 定义本地限流器限流规则
type LocalRule struct {
	Rate  int // 本地限流速率 (每秒请求数)
	Burst int // 本地限流突发容量，令牌桶的令牌数量（初始的 “突发容量” 允许短时间内用掉所有预存令牌，之后会以每秒Rate个补充令牌数）
}

// IsValid 判断是否有效
func (lr *LocalRule) IsValid() bool {
	return lr.Rate > 0 && lr.Burst > 0
}

// Metrics 记录限流相关指标
type Metrics struct {
	TotalRequests   int64 // 总请求数
	AllowedRequests int64 // 允许的请求数
	DeniedRequests  int64 // 被拒绝的请求数
	RedisErrors     int64 // Redis 错误数
}

type LocalLimiter struct {
	Limiter      *rate.Limiter
	LastUsedTime time.Time // 最新一次使用时间
}

// Limiter 基于 Redis 的分布式限流器
type Limiter struct {
	// Redis 客户端，调用redis.InitRedis()方法，或则在调用构造方法时使用WithRedisClient()设置redis客户端
	redisClient   *redis.Client
	prefix        string                   // Redis 键前缀，用于区分不同应用
	defaultRule   Rule                     // 默认限流规则
	rules         map[string]Rule          // 自定义的限流规则
	rulesMutex    sync.RWMutex             // 保护规则映射的读写锁
	localLimiters map[string]*LocalLimiter // 本地限流器缓存
	localMutex    sync.RWMutex             // 保护本地限流器的互斥锁
	luaScript     *redis.Script            // Lua 脚本获取当前窗口的请求计数，保证限流操作原子性
}

// Option 定义限流器的配置选项
type Option func(*Limiter)

// WithRedisClient 设置 Redis 连接客户端
func WithRedisClient(client *redis.Client) Option {
	return func(rl *Limiter) {
		rl.redisClient = client
	}
}

// WithPrefix 设置 Redis 键前缀
func WithPrefix(prefix string) Option {
	return func(rl *Limiter) {
		rl.prefix = prefix
	}
}

// 默认规则
var defaultRule = Rule{
	Size:        time.Minute,
	MaxRequests: 1000,
	LocalRule: LocalRule{
		Rate:  100,
		Burst: 200,
	},
}

// WithDefaultRule 设置默认限流规则
// size 滑动窗口大小
// maxRequests 窗口内最大请求数
func WithDefaultRule(size time.Duration, maxRequests int, localRules ...LocalRule) Option {
	return func(rl *Limiter) {
		rl.defaultRule = Rule{
			Size:        size,
			MaxRequests: maxRequests,
		}
		if len(localRules) > 0 {
			rl.defaultRule.LocalRule = localRules[0]
		} else {
			// 未定义则使用默认值
			rl.defaultRule.LocalRule = defaultRule.LocalRule
		}
	}
}

// NewLimiter 创建一个新的 Redis 限流器
func NewLimiter(opts ...Option) *Limiter {
	// 初始化限流器并设置默认值
	limiter := &Limiter{
		redisClient:   commonredis.Client,
		prefix:        "ratelimit:",
		rules:         make(map[string]Rule),
		localLimiters: make(map[string]*LocalLimiter),
		// 默认滑动窗口限流器规则：1分钟内最多1000个请求
		// 默认本地限流器规则：每秒允许100次请求，令牌桶有200个有效令牌
		defaultRule: defaultRule,
	}
	// 如果有调用redis.InitRedis()方法初始化redis连接客户端直接使用即可
	if commonredis.Client != nil {
		limiter.redisClient = commonredis.Client
	}

	// 应用配置选项
	for _, opt := range opts {
		opt(limiter)
	}

	// 初始化 Lua 脚本，获取当前窗口的请求计数实现滑动窗口限流算法
	limiter.luaScript = redis.NewScript(`
		-- 生成唯一的请求ID（每次请求进来都随机生成）
		local requestId = ARGV[1]
		-- 当前时间戳（毫秒）
		local now = tonumber(ARGV[2])
		-- 窗口大小（毫秒）
		local windowSize = tonumber(ARGV[3])
		-- 窗口起始时间
		local windowStart = now - windowSize
		-- 最大请求数
		local maxRequests = tonumber(ARGV[4])
		
		-- 清理窗口外的请求记录
		-- ZREMRANGEBYSCORE：Redis 有序集合的分数范围删除命令
		-- ZREMRANGEBYSCORE key 分数下限 分数上限
		redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, windowStart)
		
		-- 检查请求是否已存在（避免重复计数）
		-- ZSCORE：Redis 有序集合的分数查询命令
		-- ZSCORE key 成员
		local exists = redis.call('ZSCORE', KEYS[1], requestId)
		if exists then
			-- 已存在的请求视为允许（返回1）
			return 1
		end
		
		-- 获取当前窗口内的请求数
		local currentCount = redis.call('ZCARD', KEYS[1])
		
		-- 判断是否允许请求
		if currentCount < maxRequests then
			-- 添加当前请求
			redis.call('ZADD', KEYS[1], now, requestId)
			-- 设置过期时间（2倍窗口大小），防止内存泄漏
			-- 2倍窗口大小，确保历史数据不被意外删除
			redis.call('EXPIRE', KEYS[1], math.ceil(windowSize / 1000 * 2))
			return 1  -- 允许请求
		else
			return 0  -- 拒绝请求
		end
	`)

	// 启动定时任务清理过期的本地限流器
	go limiter.startClearExpireLocalLimiters()

	return limiter
}

// startClearExpireLocalLimiters 启动定时任务清理过期的本地限流器，防止内存泄漏
// 超过一小时未使用则清理掉
func (rl *Limiter) startClearExpireLocalLimiters() {
	ticker := time.NewTicker(time.Hour)
	defer ticker.Stop()
	for range ticker.C {
		rl.localMutex.Lock()
		for key, el := range rl.localLimiters {
			if time.Since(el.LastUsedTime) > time.Hour {
				delete(rl.localLimiters, key)
			}
		}
		rl.localMutex.Unlock()
	}
}

// AddRule 添加限流规则（本地规则不添加则默认不启用本地限流器）
// size 滑动窗口大小
// maxRequests 在滑动窗口大小的时间内允许请求的数量
// localRules 本地限流器配置
func (rl *Limiter) AddRule(name string, size time.Duration, maxRequests int, localRules ...LocalRule) {
	rl.rulesMutex.Lock()
	defer rl.rulesMutex.Unlock()
	rule := Rule{
		Name:        name,
		Size:        size,
		MaxRequests: maxRequests,
	}
	if len(localRules) > 0 {
		rule.LocalRule = localRules[0]
	}
	rl.rules[name] = rule
}

// AddRules 批量添加限流规则（本地规则不添加则默认不启用本地限流器）
func (rl *Limiter) AddRules(rules []Rule) {
	rl.rulesMutex.Lock()
	defer rl.rulesMutex.Unlock()
	for _, rule := range rules {
		rl.rules[rule.Name] = rule
	}
}

// getRule 获取适用的限流规则，优先级：自定义规则 > 默认规则
func (rl *Limiter) getRule(name string) Rule {
	rl.rulesMutex.RLock()
	defer rl.rulesMutex.RUnlock()

	// 检查自定义规则
	if rule, ok := rl.rules[name]; ok {
		return rule
	}

	// 使用默认规则
	return rl.defaultRule
}

// generateRequestID 生成唯一请求ID，防止重复计数
func (rl *Limiter) generateRequestID() string {
	// 主方案：使用crypto/rand生成加密级别的随机数
	b := make([]byte, 16)
	_, err := rand.Read(b)
	if err == nil {
		// 使用URL安全的Base64编码，去除填充字符
		return base64.URLEncoding.WithPadding(base64.NoPadding).EncodeToString(b)
	}
	// 备选方案：当加密随机数生成失败时使用
	return util.RandomStr(16)
}

// getLocalLimiter 获取或创建本地限流器
func (rl *Limiter) getLocalLimiter(key string, localRule LocalRule) *LocalLimiter {
	rl.localMutex.RLock()
	localLimiter, exists := rl.localLimiters[key]
	rl.localMutex.RUnlock()

	if !exists {
		rl.localMutex.Lock()
		defer rl.localMutex.Unlock()

		// 双重检查，防止并发创建
		if localLimiter, exists = rl.localLimiters[key]; !exists {
			localLimiter = &LocalLimiter{
				Limiter: rate.NewLimiter(rate.Limit(localRule.Rate), localRule.Burst),
			}
			rl.localLimiters[key] = localLimiter
		}
	}

	return localLimiter
}

// Allow 检查是否允许请求通过限流
// identifiers 是可选的额外标识，用于更细粒度的限流（如用户ID、IP等）
func (rl *Limiter) Allow(ctx context.Context, name string, identifiers ...string) (bool, error) {
	// 获取适用的限流规则
	rule := rl.getRule(name)

	// 构建限流键，包含名称和额外标识
	keyParts := []string{rl.prefix, name}
	keyParts = append(keyParts, identifiers...)
	redisKey := fmt.Sprintf("%s", strings.Join(keyParts, ":"))
	if commonredis.Client != nil {
		redisKey = commonredis.FullKey(redisKey)
	}

	// 执行Lua脚本进行限流判断
	// 获取当前窗口的请求计数
	var redisAllowed int
	var err error
	if rl.redisClient != nil {
		// 生成唯一请求ID，防止重复计数
		requestID := rl.generateRequestID()
		// 当前时间戳（毫秒）
		now := time.Now().UnixMilli()
		redisAllowed, err = rl.luaScript.Run(
			ctx, rl.redisClient, []string{redisKey},
			requestID, now, rule.Size.Milliseconds(), rule.MaxRequests,
		).Int()
	} else {
		err = fmt.Errorf("未设置redis客户端，或者未执行redis.InitRedis初始化redis客户端连接！")
	}
	// 处理Redis错误
	if err != nil {
		// 如果启用了本地限流作为降级方案，则使用本地限流
		var localRule LocalRule
		if rule.LocalRule.IsValid() {
			localRule = rule.LocalRule
		} else {
			localRule = rl.defaultRule.LocalRule
		}
		fmt.Printf("123|%v\r\n", localRule)
		if localRule.IsValid() {
			localLimiter := rl.getLocalLimiter(redisKey, localRule)
			allowed := localLimiter.Limiter.Allow()
			localLimiter.LastUsedTime = time.Now()
			return allowed, fmt.Errorf("redis error: %w 改用本地速率限制器", err)
		}

		// 未启用本地限流，返回错误
		return false, fmt.Errorf("redis error: %w", err)
	}

	// 判断是否允许请求
	allowed := redisAllowed == 1
	return allowed, nil
}
